// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <vector>

#include <tbir/ir/stmt_functor.h>

namespace tbir::ir {

    class YieldDetector : public StmtExprVisitor {
    public:
        const std::vector<HLOYield> &GetYields(const BaseFunc &f) {
            yields_.clear();
            VisitExpr(f);
            return yields_;
        }

        void VisitStmt_(const HLOYieldNode *op) override {
            yields_.push_back(runtime::GetRef<HLOYield>(op));
        }

    private:
        std::vector<HLOYield> yields_;
    };

    class YieldLabelMutator : public StmtExprMutator {
    public:
        BaseFunc MutateFunc(const BaseFunc &f) {
            return runtime::Downcast<BaseFunc>(StmtExprMutator::VisitExpr(f));
        }

        Stmt VisitStmt_(const HLOYieldNode *op) override {
            auto node = CopyOnWrite(op);
            node->label = IntImm(runtime::DataType::Int(64), yield_index++);
            return HLOYield(node);
        }

    private:
        int64_t yield_index = 1;
    };

}  // namespace tbir::ir
